import argparse, os, sys, threading, time
from gevent.event import Event
from locust import HttpUser, task, between, events
from locust.exception import StopUser
from config.all_paths import get_root_path
from utils.file_handling.excel_handling import ExcelUtil
from utils.interface_auth.auth_tool import Authentication
from utils.performance_tool.config_loader import load_config, get_env
# from utils.performance_tool.load_shapes import create_load_shape
from utils.request_tool.send_request import RequestSender
from config.global_variable import ENV_DICT


# ====================== 命令行参数解析 ======================
def parse_custom_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--env', required=True, help='目标环境名称')
    parser.add_argument('--role', required=True, help='测试角色类型')
    parser.add_argument('--profile_name', required=True, help='测试场景名称')
    args, remaining = parser.parse_known_args()
    sys.argv = sys.argv[:1] + remaining
    return args

args = parse_custom_args()
ENV = args.env
ROLE = args.role
PROFILE_NAME = args.profile_name

print(f"✅ 解析参数成功: ENV={ENV}, ROLE={ROLE}, PROFILE_NAME={PROFILE_NAME}")

# ====================== 配置加载 ======================
config = load_config()
get_env(ENV, ROLE)
profile_config = config['profiles'].get(PROFILE_NAME, {})
print(f"最终读取的配置参数:{profile_config}")

if not profile_config:
    print(f"❌ 错误: 未找到指定的配置文件或指定的profile_name({PROFILE_NAME})对应的配置")
    sys.exit(1)




# ====================== 读取测试数据 ======================
excelName = 'data/case_excel/wutong/uat/压测接口示例.xlsx'
testData = ExcelUtil(os.path.join(get_root_path(), excelName)).read_excel()

# ====================== 日志与统计 ======================
results = []

@events.request.add_listener
def on_request(request_type, name, response_time, response_length, **kwargs):
    results.append({
        "type": request_type,
        "name": name,
        "response_time": response_time,
        "response_length": response_length,
        "timestamp": time.time()
    })

@events.test_start.add_listener
def on_test_start(environment, **kwargs):
    print(f"\n=== 启动 {PROFILE_NAME} 场景测试 ===")
    print(f"目标环境: {ENV}")
    print(f"测试角色: {ROLE}")
    print(f"基准URL: {ENV_DICT.get('url')}")

@events.test_stop.add_listener
def on_test_stop(environment, **kwargs):
    print("\n=== 测试结束，统计结果如下 ===")
    stats = environment.stats.total
    print(f"请求总数: {stats.num_requests}")
    print(f"失败数: {stats.num_failures}")
    print(f"失败率: {stats.fail_ratio:.2%}")
    print(f"平均响应时间: {stats.avg_response_time:.2f} ms")
    print(f"最大响应时间: {stats.max_response_time:.2f} ms")
    print(f"中位响应时间: {stats.median_response_time:.2f} ms")
    print(f"TPS: {stats.total_rps:.2f}")


# # ========== 启用 LoadShape ==========
# load_test_shape = create_load_shape(profile_config)
# print(f"负载控制器已创建: {load_test_shape}")

# ====================== 基础用户类 ======================
user_counter = 0
user_lock = threading.Lock()

class BaseUser(HttpUser):
    wait_time = between(0.5, 2)
    # load_shape = load_test_shape
    abstract = True

    def on_start(self):
        global user_counter
        with user_lock:
            self.user_id = user_counter
            user_counter += 1

            # 调用负载控制器的 tick 方法来获取当前阶段的用户数量和生成速率
        # self.update_load_shape()

        auth = Authentication(
            ENV_DICT.get('username'),
            ENV_DICT.get('password'),
            ENV_DICT.get('url'),
        )

        auth_method_map = {
            "test_wt": "auth_wutong",
            "test_th": "auth_th",
            "test_inter": "auth_inter"
        }

        auth_method_name = auth_method_map.get(ENV, "auth_wutong")
        if hasattr(auth, auth_method_name):
            self.s = getattr(auth, auth_method_name)()
            print(f"用户 {self.user_id} 使用 {auth_method_name} 完成认证")
            print(f"接口鉴权后的head是:",self.s.headers)
        else:
            raise ValueError(f"认证方法 {auth_method_name} 不存在")

    # def update_load_shape(self):
    #     shape_result = self.load_shape.tick()
    #     if shape_result:
    #         users, spawn_rate = shape_result
    #         self.users = users
    #         self.spawn_rate = spawn_rate
    #         print(f"负载控制器返回: 当前用户数 = {self.users}, 每秒生成速率 = {self.spawn_rate}")
    #     else:
    #         print("负载控制器已完成所有阶段")

# ====================== 各种测试用户类 ======================
class BaselineUser(BaseUser):
    @task
    def baseline_task(self):
        print("======使用的是baseline方法")
        result = RequestSender().send_request(self.s, testData[0]).json()
        print("接口请求结果:", result)

class StressUser(BaseUser):
    @task(5)
    def stress_task(self):
        print("======使用的是Stress方法")
        RequestSender().send_request(self.s, testData[0])

class WaveUser(BaseUser):
    @task
    def wave_task(self):
        print("======使用的是wave方法")
        RequestSender().send_request(self.s, testData[0])

# ========== 并发用户（集合点机制） ==========
expected_users = profile_config.get("users", 1)
arrival_count = 0
arrival_lock = threading.Lock()
start_event = Event()

class ConcurrencyUser(BaseUser):
    def on_start(self):
        super().on_start()
        global arrival_count
        with arrival_lock:
            arrival_count += 1
            print(f"[集合] 用户 {self.user_id} 抵达（{arrival_count}/{expected_users}）")
            if arrival_count >= expected_users:
                print("[集合] 所有用户到齐，开始统一释放！")
                start_event.set()

    @task
    def concurrency_task(self):
        print("======使用的是concurrency方法")
        try:
            start_event.wait(timeout=30)
            print(f"[并发] 用户 {self.user_id} 发起请求")
            RequestSender().send_request(self.s, testData[0])
            raise StopUser()
        except Exception as e:
            print(f"[并发异常] 用户 {self.user_id} 错误：{e}")
            raise StopUser()

# ====================== 动态绑定场景类 ======================
class DynamicTestClass(eval(profile_config.get("user_class", "BaselineUser"))):
    pass

class TestScenario(DynamicTestClass):
    pass

# def get_shape_class():
#     return create_load_shape(profile_config)



def get_shape_class():
    from utils.performance_tool.load_shapes import create_load_shape
    # 确保懒加载，在运行时创建负载控制器
    return create_load_shape(profile_config)


shape_class = get_shape_class()
print(f"✅ shape_class 被创建: {shape_class}, 类型: {type(shape_class)}")


# shape_class = DebugShape()


### 命令行
# locust -f locustfile.py --env test_wt --role doctor --profile_name baseline --run-time 1m --host=https://wt-stg.zenggao.net  恒定负载
# locust -f locustfile.py --env test_wt --role doctor --profile_name incremental --run-time 2m --host=https://wt-stg.zenggao.net  阶梯递增
# locust -f locustfile.py --env test_wt --role doctor --profile_name peak --run-time 1m --host=https://wt-stg.zenggao.net  峰值压力
# locust -f locustfile.py --env test_wt --role doctor --profile_name concurrency --run-time 30s --host=https://wt-stg.zenggao.net  并发集合点
# locust -f locustfile.py --env test_wt --role doctor --profile_name wave --run-time 3m --host=https://wt-stg.zenggao.net  波浪式压力